import pandas as pd
import matplotlib.pyplot as plt
import sys
from services.DayKlineService import *
kline = DayKlineService()

pd.set_option('display.max_columns', None) # 展示所有列

alldate = kline.getAllData(sys.argv[1])
stock_data = pd.DataFrame(alldate)
stock_data['date'] = pd.to_datetime(stock_data['trade_date'])
# 求该股票每日涨跌幅
stock_data['change_pct'] = stock_data['close'].pct_change()
df = stock_data[['date', 'open', 'close', 'high', 'low', 'change_pct']]
df['date'] = pd.to_datetime(df['date'])  # 将交易日期字符串变为日期类型
print(df.head())
# 计算均线
ma_short = 5  #短期均线，ma代表：moving_average
ma_long = 20  #长期均线，ma代表：moving_average
df['ma_short'] = df['close'].rolling(ma_short).mean()
df['ma_long'] = df['close'].rolling(ma_long).mean()
print(df.head(10))
#补全上面均线缺失值：补全方式采用扩展窗口函数expanding，移动计算前面所有值之和的均值
df['ma_short'].fillna(value=df['close'].expanding().mean(),inplace=True)
df['ma_long'].fillna(value=df['close'].expanding().mean(),inplace=True)
print(df.head(10))
# 找买入信号
# 当天的短期均线大于等于长期均线
condition1 = (df['ma_short'] >= df['ma_long'])
# 上一个交易日的短期均线小于长期均线
condition2 = (df['ma_short'].shift(1) < df['ma_long'].shift(1))
# 将买入信号当天的signal设为1
df.loc[condition1 & condition2, 'signal'] = 1

# 找卖出信号
# 当天的短期均线小于等于长期均线
condition1 = (df['ma_short'] <= df['ma_long'])
# 上一个交易日的短期均线大于长期均线
condition2 = (df['ma_short'].shift(1) > df['ma_long'].shift(1))
# 将买入信号当天的signal设为0
df.loc[condition1 & condition2,'signal'] = 0

# 浏览产生交易信号的日期
print(df[df['signal'].notnull()].head(10))

# 新建列Pos表示仓位，将交易信号下移一格，表示第二天买入，1表示满仓，0表示空仓
df['pos'] = df['signal'].shift()
# 向上填充，将买入之后的pos全部设置为1
df['pos'].fillna(method='ffill', inplace=True)
# 没有买入的pos全部设置为0
df['pos'].fillna(value=0,inplace=True)
# 预览仓位数据
print(df[['date','signal','pos']].head(20))

# 找出开盘涨停的日期：即今天的开盘价相对于昨天的收盘价上涨了9.7%以上，此处不用10%是因为由于4舍5入，涨停不一定就是10%
cond_cannot_buy = df['open'] > df['close'].shift(1) * 1.097

# 将开盘涨停且当前position为1时的'pos'设为空值
df.loc[cond_cannot_buy & (df['pos'] == 1),'pos'] = None

# 找出开盘跌停的日期，即今天的开盘价相对于昨天的收盘价跌了9.7%（1-0.097=0.903）
cond_cannot_sell = df['open'] < df['close'].shift(1) *0.903

# 将开盘跌停且当前position为0时的'pos'设为空值
df.loc[cond_cannot_sell & (df['pos'] == 0),'pos'] = None

# position为空的日期表示不能买卖。position仓位只能和前一个交易日保持一致
df['pos'].fillna(method='ffill', inplace=True)

# 资金曲线：假设起始资金为100万元的每天资金变化情况
# 首先计算资金曲线每天的涨跌幅，‘equity_change’表示资金每天的涨跌幅
# 当天空仓时，pos为0，资产涨幅为0
# 当天满仓时，pos为1，资产涨幅为股票本身的涨跌幅
df['equity_change'] = df['change_pct'] * df['pos']
# 根据每天的涨跌幅计算资金曲线
df['equity_curve'] = 10000 * (df['equity_change'] + 1).cumprod()
df = df[['date', 'change_pct', 'pos', 'equity_change', 'equity_curve']]
df.reset_index(inplace=True, drop=True)  # 重置索引，让他从0开始

# 绘制资金曲线
plt.plot(df['equity_curve'])
plt.show()